// Distributions
// distrib.hh
// (c) Yongjin Park, 2012
//

#ifndef DISTRIB_HH_
#define DISTRIB_HH_

#include <cstdlib>
#include <cmath>
#include <string>
#include <sstream>
#include <ostream>
#include <vector>
#include <limits>

#include <boost/lexical_cast.hpp>
#include <boost/math/special_functions/fpclassify.hpp>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/gamma.hpp>

#include "util.hh"

using namespace std;
using namespace boost;


////////////////////////////////////////////////////////////////
// latent variable or latent mass
struct latent_t
{
  explicit latent_t(double _z) : z(_z) {
#ifdef DEBUG
    assert( _z >= -1e-5 );
#endif
    if( z < 0. ) z = 0.;
  }
  double z;
};

////////////////////////////////////////////////////////////////
// learning rate
struct rate_t
{
  explicit rate_t(double _e) : eta(_e) {
#ifdef DEBUG
    assert( _e >= -1e-5 && _e <= (1.+1e-5) );
#endif
    if( eta < 0. ) eta = 0.;
    if( eta > 1. ) eta = 1.;
  }
  double eta;
};

struct count_t
{
  explicit count_t(double _c): count(_c) {
#ifdef DEBUG
    assert( _c >= -1e-5 );    
#endif
    if( count < 0. ) count = 0.;
  }
  double count;
};

////////////////////////////////////////////////////////////////
// abstract class of distribution
template<class T>
class distrib_t
{

public:

  virtual void add_stat(const T& rhs, const latent_t& latent, const rate_t& eta)= 0;
  virtual void add_stat(const T& rhs, const latent_t& latent)=0;
  virtual void add_stat(const T& rhs)=0;
  virtual void discount_stat(const rate_t& c) = 0;

  // gradient of single observation
  virtual double gradient(const T& obs) const = 0;

  // variational approximation
  virtual double update_param() = 0;

  // log likelihood based on MLE parameters
  virtual double log_lik() const = 0;

  // log marginal likelihood
  virtual double log_marg() const = 0;

  virtual size_t size() const = 0;

#ifdef DEBUG
  virtual const T& get_stat() const = 0;
  virtual double get_total () const = 0;
#endif

#ifdef DEBUG
  virtual bool optimized() const = 0;  // prevent using unoptimized parameters
#endif
};

class uni_distrib_t : public distrib_t<double>
{
public:
  virtual double mean() const = 0;
  virtual double sd() const = 0;
};

////////////////////////////////////////////////////////////////
// useful for batch update of network-related
class network_distrib_t : public distrib_t<double>
{
public:
  struct edge_t : public count_t
  {
    explicit edge_t(double e) : count_t(e) {}
  };
  struct tot_t : public count_t
  {
    explicit tot_t(double h) : count_t(h) {}
  };

  struct hyper_t
  {
    explicit hyper_t(double _a, double _b) : a(_a), b(_b) {
#ifdef DEBUG
      assert( a> 0. && b > 0. );
#endif
    }
    double a, b;
  };

public:

  network_distrib_t();
  network_distrib_t( const hyper_t& hyper );
  virtual ~network_distrib_t();

  void add_stat(const double& rhs, const latent_t& latent, const rate_t& eta);
  void add_stat(const double& rhs, const latent_t& latent);
  void add_stat(const double& rhs);
  void discount_stat(const rate_t& c);

  void add_edge(const edge_t& x);
  void add_tot(const tot_t& x);
  void add_edge(const edge_t& x, const rate_t&);
  void add_tot(const tot_t& x, const rate_t&);

  double num_edge() const;
  double num_tot() const;

  // gradient of single observation
  virtual double gradient(const edge_t& e, const tot_t& t) const = 0;

  // log predictive for locally collapsed variational inference
  virtual double log_pred(const edge_t& e, const tot_t& t) const = 0;

  // mean parameter for link prediction
  virtual double mean() const = 0;

  // distribution-wise update
  void operator += (const network_distrib_t&);

  size_t size() const;
#ifdef DEBUG
  const double& get_stat() const;
  double get_total () const;
#endif

  bool read( string );
  string write() const;


#ifdef DEBUG
  bool optimized() const;  // prevent using unoptimized parameters
#endif

protected:
  double a;
  double b;
  double X;
  double N;
  bool opt;
};

////////////////////////////////////////////////////////////////
// bernoulli-beta distribution
class bernoulli_t : public network_distrib_t
{
public:

  typedef network_distrib_t::hyper_t hyper_t;

  bernoulli_t();
  bernoulli_t(const hyper_t& hyper);
  ~bernoulli_t();

  // calculate gradient
  double gradient(const double& obs) const;
  double gradient(const edge_t& e, const tot_t& t) const;

  // log predictive
  double log_pred(const edge_t& e, const tot_t& t) const;

  // mean parameter
  double mean() const;

  // variational approximation
  double update_param();

  // log likelihood based on MLE parameters
  double log_lik() const;

  // log marginal likelihood
  double log_marg() const;

  bool read( string );

private:
  double thetaEdge;
  double thetaHole;
};


////////////////////////////////////////////////////////////////
// poisson-gamma distribution
class poisson_t : public network_distrib_t
{  
public:

  typedef network_distrib_t::hyper_t hyper_t;

  poisson_t();
  poisson_t(const hyper_t& hyper);
  ~poisson_t();

  // calculate gradient
  double gradient(const double& obs) const;
  double gradient(const edge_t& e, const tot_t& t) const;

  // log predictive
  double log_pred(const edge_t& e, const tot_t& t) const;

  // mean parameter
  double mean() const;

  // variational approximation
  double update_param();

  // log likelihood based on MLE parameters
  double log_lik() const;

  // log marginal likelihood
  double log_marg() const;

  bool read( string );

private:
  double lambda;
  double logLambda;
};






#endif /* DISTRIB_HH_ */
